import os
import pandas as pd
import torch
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
from PIL import Image

# 数据集均值和标准差，用于归一化
# 这些值是针对GTSRB数据集计算得出的，可以继续使用
STATS = ([0.3403, 0.3121, 0.3214], [0.2751, 0.2642, 0.2706])

def get_data_loaders(data_dir, batch_size=64):
    """
    创建并返回训练和测试数据加载器 (适配Kaggle数据集结构)。

    Args:
        data_dir (str): 包含Train/Test图片和CSV文件的数据集根目录。
        batch_size (int): 训练和测试的批量大小。

    Returns:
        tuple: (train_loader, test_loader)
    """
    # 训练数据转换：数据增强 + 归一化
    train_transform = transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.RandomRotation(10),
        transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
        transforms.ColorJitter(brightness=0.2, contrast=0.2),
        transforms.ToTensor(),
        transforms.Normalize(*STATS, inplace=True)
    ])

    # 测试数据转换：仅调整尺寸和归一化
    test_transform = transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor(),
        transforms.Normalize(*STATS, inplace=True)
    ])

    # Kaggle数据集结构下的路径
    train_csv_path = os.path.join(data_dir, 'Train.csv')
    test_csv_path = os.path.join(data_dir, 'Test.csv')

    train_dataset = GTSRBKaggleDataset(root_dir=data_dir, csv_file=train_csv_path, transform=train_transform)
    test_dataset = GTSRBKaggleDataset(root_dir=data_dir, csv_file=test_csv_path, transform=test_transform)

    train_loader = DataLoader(
        dataset=train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=4,
        pin_memory=True
    )

    test_loader = DataLoader(
        dataset=test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=4,
        pin_memory=True
    )

    return train_loader, test_loader

class GTSRBKaggleDataset(Dataset):
    """
    用于Kaggle版GTSRB数据集的通用Dataset类。
    通过读取CSV文件来加载训练集或测试集。
    """
    def __init__(self, root_dir, csv_file, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        
        # 从CSV文件读取标注
        self.annotations = pd.read_csv(csv_file)
        
        # 拼接完整图像路径
        self.image_paths = self.annotations['Path'].apply(lambda f: os.path.join(root_dir, f)).tolist()
        self.labels = self.annotations['ClassId'].tolist()

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        label = self.labels[idx]
        
        try:
            image = Image.open(img_path).convert("RGB")
        except FileNotFoundError:
            print(f"错误：找不到文件 {img_path}")
            # 返回一个占位符图像和标签，或者可以引发异常
            return torch.zeros(3, 32, 32), -1 
        
        if self.transform:
            image = self.transform(image)
            
        return image, label

if __name__ == '__main__':
    # --- 测试数据加载器 (Kaggle版) ---
    # 请将 'path/to/your/gtsrb' 替换为你的Kaggle版GTSRB数据集根目录
    # 这个目录下应该直接包含 Train.csv, Test.csv, 以及 Train/和Test/ 文件夹
    
    # 这是一个示例路径，你需要修改它
    kaggle_gtsrb_path = './data' 
    
    if not os.path.exists(os.path.join(kaggle_gtsrb_path, 'Train.csv')):
         print("="*60)
         print(f"错误: 无法在 '{os.path.abspath(kaggle_gtsrb_path)}' 找到 'Train.csv'。")
         print("请确认您下载的Kaggle数据集已解压，并将包含CSV文件的根目录路径正确设置。")
         print("例如, 如果你的路径是 './data'，那么该目录下应有 'Train.csv', 'Test.csv' 等。")
         print("="*60)
    else:
        print("正在测试Kaggle版数据加载器...")
        # 注意：这里我们使用 data_dir 指向包含 Train.csv 和 Test.csv 的目录
        train_loader, test_loader = get_data_loaders(kaggle_gtsrb_path, batch_size=4)

        # 从训练加载器中取一个批次
        train_images, train_labels = next(iter(train_loader))
        print(f"\n成功从 train_loader 加载数据！")
        print(f"训练图像批次尺寸: {train_images.shape}")
        print(f"训练标签批次尺寸: {train_labels.shape}")
        print(f"标签示例: {train_labels[:4]}")

        # 从测试加载器中取一个批次
        test_images, test_labels = next(iter(test_loader))
        print(f"\n成功从 test_loader 加载数据！")
        print(f"测试图像批次尺寸: {test_images.shape}")
        print(f"测试标签批次尺寸: {test_labels.shape}")
        print(f"标签示例: {test_labels[:4]}")
        print("\nKaggle版数据加载器测试通过！") 